import torch
import torchvision
import torchvision.transforms as transforms
from vit import write_json
import os
import random
import json

def get_transform():
    return transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])


if __name__ == '__main__':
    dir_data = "/home/leslie/project/ViT/dataset/肉类_30"
    model = torch.load('model_test_best.pth')
    transform = get_transform()

    data = torchvision.datasets.ImageFolder(dir_data, transform=transform)
    label_class = {value: key for key, value in data.class_to_idx.items()}
    #print(label_class)

    pred = {}
    for i in range(len(data)):
    # for i in range(100):
        try:
            a, b = data.__getitem__(random.randint(0,len(data)))
            # a, b = data.__getitem__(i)
            a.unsqueeze_(0)
            a = a.to("cuda:0")
            res = model(a)
            m = torch.argmax(res).item()
            if label_class[b] not in pred.keys():
                pred[label_class[b]] = [m]
            else:
                pred[label_class[b]].append(m)
            if i % 200 == 0:
                print("processed {} images".format(i))

        except Exception as e:
            write_json("re.json", pred, os.getcwd())
            print(e)
    write_json("re.json", pred, os.getcwd())



